import requests
import logging
import time


class FlowerManager(object):
    """
    通过flower来管理celery
    """

    def __init__(self, domain):
        """
        :param domain: flower域名
        """
        if domain[-1] == "/":
            self.domain = domain[0:-1]
        else:
            self.domain = domain

    def __flower_requests_withjson__(self, method, url, **kwargs):
        """封装请求"""
        try:
            response = requests.request(method, url, **kwargs)
            try:
                return response.json()
            except Exception as e:
                raise Exception(
                    f"【 {time.asctime()} 】\nFlower api response exception 【 {url} 】【{response.status_code}】--> 【 {response.text} 】 ")
        except Exception as e:
            raise Exception(f"【 {time.asctime()} 】\nFlower api requests exception 【 {url} 】")

    def get_workers(self):
        """workers清单"""
        url = self.domain + "/api/workers"
        return self.__flower_requests_withjson__(method="get", url=url)

    def shutdown_worker(self, worker):
        """关闭worker"""
        url = self.domain + f"/api/worker/shutdown/{worker}"
        return self.__flower_requests_withjson__(method="post", url=url)

    def restart_worker_pool(self, worker):
        """重启worker池"""
        url = self.domain + f"/api/worker/pool/restart/{worker}"
        return self.__flower_requests_withjson__(method="post", url=url)

    def grow_worker_pool(self, worker, n: int):
        """增大worker池"""
        url = self.domain + f"/api/worker/pool/grow/{worker}?n={n}"
        return self.__flower_requests_withjson__(method="post", url=url)

    def shrink_worker_pool(self, worker, n: int):
        """减少worker池"""
        url = self.domain + f"/api/worker/pool/grow/{worker}?n={n}"
        return self.__flower_requests_withjson__(method="post", url=url)

    def autoscale_worker_pool(self, worker, min: int = 3, max: int = 10):
        """自动设置worker池"""
        url = self.domain + f"/api/worker/pool/grow/{worker}?min={min}&max={max}"
        return self.__flower_requests_withjson__(method="post", url=url)

    def add_consumer_worker_queue(self, worker, queue):
        """为worker执行消费队列"""
        url = self.domain + f"/api/worker/queue/add-consumer/{worker}?queue={queue}"
        return self.__flower_requests_withjson__(method="post", url=url)

    def cancle_consumer_worker_queue(self, worker, queue):
        """为worker停止消费队列"""
        url = self.domain + f"/api/worker/queue/cancel-consumer/{worker}?queue={queue}"
        return self.__flower_requests_withjson__(method="post", url=url)

    def get_tasks(self):
        """获取任务清单"""
        url = self.domain + f"/api/tasks"
        return self.__flower_requests_withjson__(method="get", url=url)

    def get_task_types(self):
        """获取任务类型"""
        url = self.domain + f"/api/task/types"
        return self.__flower_requests_withjson__(method="get", url=url)

    def get_task_info(self, task_id):
        """获取一个任务详细信息"""
        url = self.domain + f"/api/task/info/{task_id}"
        return self.__flower_requests_withjson__(method="get", url=url)

    def async_apply_task(self, task_type, args=None):
        """异步执行任务(celery会检查方法是否存在)"""
        url = self.domain + f"/api/task/async-apply/{task_type}"
        data = {
            "args": args
        }
        return self.__flower_requests_withjson__(method="post", url=url, json=data)

    def send_task(self, task_type, args=None):
        """异步执行任务（celery不会检查任务是否存在）"""
        url = self.domain + f"/api/task/send-task/{task_type}"
        data = {
            "args": args
        }
        return self.__flower_requests_withjson__(method="post", url=url, json=data)

    def get_task_result(self, task_id):
        """获取任务的执行结果"""
        url = self.domain + f"/api/task/result/{task_id}"
        return self.__flower_requests_withjson__(method="get", url=url)

    def timeout_task(self, worker, task_type, soft_timeout: int = 30, hard_timeout: int = 100):
        """为worker某一种任务设置软硬时限"""
        url = self.domain + f"/api/task/timeout/{task_type}" + f"?soft={soft_timeout}&hard={hard_timeout}&workername={worker}"
        return self.__flower_requests_withjson__(method="post", url=url)

    def rate_limit_task(self, worker, task_type, ratelimit: int = 200):
        """为worker某种任务进行限速"""
        url = self.domain + f"/api/task/rate-limit/{task_type}" + f"ratelimit={ratelimit}&workername={worker}"
        return self.__flower_requests_withjson__(method="post", url=url)

    def revoke_task(self, task_id, signal="SIGUSR1"):
        """终止一个任务"""
        url = self.domain + f"/api/task/revoke/{task_id}?terminate=true&signal={signal}"
        return self.__flower_requests_withjson__(method="post", url=url)


def test():
    flower = FlowerManager(domain="http://test:test@127.0.0.1:5555")
    res = flower.get_workers()
    # res = flower.async_apply_task("celery_tasks.test.tasks.test_task", args=(1, 4))
    # print(res)
    # res = flower.timeout_task(worker="celery@rockontroldeMBP", task_type="celery_tasks.test.tasks.test_task")
    # res = flower.revoke_task(task_id=res.get("task-id"))
    # res = flower.revoke_task(task_id="9986edf2-e7d5-4285-9c4f-62fbb348eb96")
    print(res)


if __name__ == '__main__':
    test()
